#define DEBUG_TYPE "abcd"
#include "llvm/Pass.h"
#include "llvm/CallGraphSCCPass.h"
#include "llvm/Function.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Module.h"
#include "llvm/Constants.h"
#include "llvm/Instructions.h"
#include "llvm/Support/IRBuilder.h"
#include "llvm/Pass.h"
#include "llvm/Function.h"
#include "llvm/InstrTypes.h"
#include "llvm/Instructions.h"
#include "llvm/Analysis/Dominators.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Transforms/Utils/PromoteMemToReg.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Support/raw_ostream.h"

#include "InequalityGraph.h"

#include <set>
#include <algorithm>
#include <stack>
#include <iostream>
#include <vector>

using namespace llvm;

using std::cout;
using std::set;
using std::map;
using std::vector;
using std::pair;
using std::stack;
using std::string;

namespace llvm {
STATISTIC(ABCFound, "Array Bound Check Found");
STATISTIC(ABCD, "Redundant Array Bound Check Eliminated");

template<class T>
void PRINT(T* x, const std::string& str = "") {
	if (MY_DEBUG) {
		errs() << str << ":\n\t";
		x->print(errs());
		errs() << '\n';
	}
}

bool isPINode(PHINode* phiNode) {
	return phiNode->getNumIncomingValues() == 1;
}

PHINode* getPiDef(Value* val, BasicBlock* body) {
	for (BasicBlock::iterator i = body->begin(), e = body->end(); i != e; ++i) {
		if (PHINode * phiNode = dyn_cast<PHINode>(&(*i))) {
			if (!isPINode(phiNode))
				continue;
			Value* piArg = phiNode->getIncomingValue(0);
			if (piArg == val)
				return phiNode;
		}
	}
	return NULL;
}

std::pair<bool, int> isConstInt(Value* val) {
	if (ConstantInt * constInt = dyn_cast<ConstantInt>(val)) {
		PRINT(constInt, "Constant Int");

		int val = constInt->getValue().getLimitedValue();

		return std::make_pair(true, val);
	}
	return std::make_pair(false, 0);
}

struct MyFunctionPass: public FunctionPass {
	static char ID;
	MyFunctionPass() :
			FunctionPass(ID) {
	}

	// The inequality graph.
	InequalityGraph ig;

  virtual bool doInitialization(Module &M) {
    ABCFound = 0;
    ABCD = 0;
    return false;
  }

  virtual bool doFinalization(Module &M) {
    errs() << "[ABCRemove]>>>Total number of ABC checks found: "
        << ABCFound << '\n';
    errs() << "[ABCRemove]>>>Total number of ABC checks removed: "
        << ABCD << '\n';
    return false;
  }

	virtual bool runOnFunction(Function &F) {
		for (Function::iterator fb = F.begin(), fe = F.end(); fb != fe; ++fb) {
			for (BasicBlock::iterator bb = fb->begin(), be = fb->end();
					bb != be; ++bb) {
				Instruction* inst = &(*bb);

				// Handle control flows (phi nodes).
				if (PHINode * phiNode = dyn_cast<PHINode>(inst)) {
					for (int i = 0, s = phiNode->getNumIncomingValues(); i < s;
							++i) {
						PRINT(phiNode->getIncomingValue(i));
						ig.addEdge(phiNode->getIncomingValue(i), phiNode, 0);
					}
					//phiNode->getIncomingValue(1)->print(errs());
//					errs() << '\n';
					//cout << "Detected a PHI node!\n";
				}

				// Handle C3: plus and sub
				else if (BinaryOperator * binOp = dyn_cast<BinaryOperator>(
						inst)) {
					Value* lhs = binOp->getOperand(0); //->print(errs());
					Value* rhs = binOp->getOperand(1);

					Instruction::BinaryOps opT = binOp->getOpcode();
					if (opT == Instruction::Add) {
						PRINT(binOp, "Add Op");

						std::pair<bool, int> res = isConstInt(lhs);
						// If the lhs is a constant
						if (res.first) {
							ig.addEdge(rhs, binOp, res.second);
              continue;
            }
						else {
							res = isConstInt(rhs);
							if (res.first) {
								ig.addEdge(lhs, binOp, res.second);
                continue;
              }
						}

            //neither lhs nor rhs is const
            //add a new AddVertex
            AddVertex* Add = new AddVertex(binOp, lhs, rhs);
            ig.addVertex(Add);

#if 0
              errs() << "add:";
              binOp->print(errs());
              errs() << '\n';
              errs() << "lhs: ";
              lhs->print(errs());
              errs() << '\n';
              errs() << "rhs: ";
              rhs->print(errs());
              errs() << '\n';

              errs() << "added vertex: ";
              ig.vertices[binOp]->value->print(errs());
              errs() << '\n';

						errs() << "\n\n";
#endif
					} // end Add
          else if (opT == Instruction::Sub) {
						PRINT(binOp, "Sub Op");

						std::pair<bool, int> res = isConstInt(rhs);
						// If the lhs is a constant
						if (res.first)
							ig.addEdge(lhs, binOp, -res.second);

//						errs() << "\n\n";
					} // end Sub

          else if (opT == Instruction::Mul) {

            //if one of lhs or rhs is constant
						std::pair<bool, int> res = isConstInt(lhs);
						// If the lhs is a constant
						if (res.first) {
              MulVertex* Mul = new MulVertex(binOp, lhs, rhs);
              ig.addVertex(Mul);

#if 0
              errs() << "mul:";
              binOp->print(errs());
              errs() << '\n';
              errs() << "lhs: ";
              lhs->print(errs());
              errs() << '\n';
              errs() << "rhs: ";
              rhs->print(errs());
              errs() << '\n';

              errs() << "added vertex: ";
              ig.vertices[binOp]->value->print(errs());
              errs() << '\n';
#endif
            } // end lhs
						else {
							res = isConstInt(rhs);
							if (res.first) {
                MulVertex* Mul = new MulVertex(binOp, rhs, lhs);
                ig.addVertex(Mul);
              }
						} // end rhs

          } // end Mul
				} // end binOp

				// Handle C4: Add edges for predicates.
				if (BranchInst * BI = dyn_cast<BranchInst>(inst)) {
					if (BI->isConditional()) {
						if (CmpInst * CI = dyn_cast<CmpInst>(
								BI->getCondition())) {
							BasicBlock *TB = BI->getSuccessor(0); // True block
							BasicBlock *FB = BI->getSuccessor(1); // False block

							PHINode* lhsT = getPiDef(CI->getOperand(0), TB);
							PHINode* rhsT = getPiDef(CI->getOperand(1), TB);

							PHINode* lhsF = getPiDef(CI->getOperand(0), FB);
							PHINode* rhsF = getPiDef(CI->getOperand(1), FB);
							//One or two of the operands could be constants
							ConstantInt* lhsC = dyn_cast<ConstantInt>(
									CI->getOperand(0));
							ConstantInt* rhsC = dyn_cast<ConstantInt>(
									CI->getOperand(1));
							// Add edges for PI nodes.
							if (lhsT)
								ig.addEdge(lhsT->getIncomingValue(0), lhsT, 0);
							if (rhsT)
								ig.addEdge(rhsT->getIncomingValue(0), rhsT, 0);
							if (lhsF)
								ig.addEdge(lhsF->getIncomingValue(0), lhsF, 0);
							if (rhsF)
								ig.addEdge(rhsF->getIncomingValue(0), rhsF, 0);

							CmpInst::Predicate pred = CI->getPredicate();
							switch (pred) {
							case CmpInst::ICMP_UGT:
							case CmpInst::ICMP_SGT:
								//a > b  => a >= b + 1  =>  b - a <= -1
								ig.addEdge(lhsT, rhsT, -1);
								ig.addEdge(rhsF, lhsF, 0);

								ig.addEdge(lhsT, rhsC, -1);
								ig.addEdge(lhsC, rhsT, -1);

								ig.addEdge(rhsC, lhsF, 0);
								ig.addEdge(rhsF, lhsC, 0);
								break;

							case CmpInst::ICMP_UGE:
							case CmpInst::ICMP_SGE:
								ig.addEdge(lhsT, rhsT, 0);
								ig.addEdge(rhsF, lhsF, -1);

								ig.addEdge(lhsC, rhsT, 0);
								ig.addEdge(lhsT, rhsC, 0);

								ig.addEdge(rhsC, lhsF, -1);
								ig.addEdge(rhsF, lhsC, -1);
								break;

							case CmpInst::ICMP_ULT:
							case CmpInst::ICMP_SLT:
								ig.addEdge(rhsT, lhsT, -1);
								ig.addEdge(lhsF, rhsF, 0);

								ig.addEdge(rhsC, lhsT, -1);
								ig.addEdge(rhsT, lhsC, -1);

								ig.addEdge(lhsC, rhsF, 0);
								ig.addEdge(lhsF, rhsC, 0);
								break;

							case CmpInst::ICMP_ULE:
							case CmpInst::ICMP_SLE:
								ig.addEdge(rhsT, lhsT, 0);
								ig.addEdge(lhsF, rhsF, -1);

								ig.addEdge(rhsC, lhsT, 0);
								ig.addEdge(rhsT, lhsC, 0);

								ig.addEdge(lhsC, rhsF, -1);
								ig.addEdge(lhsF, rhsC, -1);
								break;

							default:
								break;
							}
						}
					}
				}

        //Handle Sign extend
				if (SExtInst * SI = dyn_cast<SExtInst>(inst)) {
					Value* op = SI->getOperand(0);
					ig.addEdge(op, SI, 0);
				}

        //Handle Swtich
        if (SwitchInst* SI = dyn_cast<SwitchInst>(inst)) {
          //not enough information for default case
          for (unsigned int i = 1; i < SI->getNumCases(); ++i) {
            BasicBlock* Succ = SI->getSuccessor(i);
            PHINode* Pi = getPiDef(SI->getCondition(), Succ);
            ig.addEdge(SI->getCaseValue(i), Pi, 0);
          }
        } //end if swtich

			}
		}
		return removeRABC(F);
	}
	bool removeRABC(Function& F) {
		bool ifChanged = false;
		std::vector<BasicBlock*> bbToMerge;
		std::vector<BasicBlock*> bbToErase;
		for (Function::iterator iter = F.begin(); iter != F.end(); iter++) {
			TerminatorInst* term = iter->getTerminator();
			BranchInst* branch = dyn_cast<BranchInst>(term);
			if (0 != branch) { //Successfully casted to a branch inst
				if (branch->isConditional()) {
					BasicBlock* fb = branch->getSuccessor(1);
					if (fb->getName().startswith("ABC_Check_Fail")) {
						if (MY_DEBUG) {
							errs() << "Array bound check found!" << "\n\r";
							errs()
									<< "Check whether if the array bound check is redundant."
									<< "\n\r";
						}
						ABCFound++;
						CmpInst *cmp = dyn_cast<CmpInst>(
								branch->getCondition());
						bool result = true;
						if (cmp == NULL) {
							// if it is a constant result, e.g. true or false
							if (ConstantInt * cmp_result =
									dyn_cast<ConstantInt>(
											branch->getCondition())) {
								//if true
								if (!(cmp_result->isOne()))
									continue;
							}
							else continue;
						}else{
						Value* index = cmp->getOperand(0);
						Value* bound = cmp->getOperand(1);
						if(MY_DEBUG)
						errs() << "index: " << *index << "\t" << "bound: "
								<< *bound << "\n\r";
						ConstantInt* conIndex = dyn_cast<ConstantInt>(index);
						ConstantInt* conBound = dyn_cast<ConstantInt>(bound);
						if (conIndex && conBound) {
							result = conIndex->getValue().slt(
									conBound->getValue());
						} else
							result = ig.demandProve(index, bound);
						}
						if (result) {
							ifChanged = true;
							if (MY_DEBUG)
								errs() << "Redundant Array Bound Check found."
										<< "\n\r";
							//Some transformation on the basic block to prepare for the merge of basic blocks
							BasicBlock* sb = branch->getSuccessor(0);
							BranchInst* ucJmp = BranchInst::Create(sb);
							ReplaceInstWithInst(branch, ucJmp);
							if(cmp)
							cmp->eraseFromParent();
							bbToMerge.push_back(sb);
							bbToErase.push_back(fb);
						} else {
							if (MY_DEBUG)
								errs() << "Valid Array Bound Check." << "\n\r";
						}
					}
				}
			}
		}
		bool ifAllMerged = true;
		for (std::vector<BasicBlock*>::iterator iter = bbToMerge.begin();
				iter != bbToMerge.end(); iter++) {
			if (!MergeBlockIntoPredecessor(*iter))
				ifAllMerged = false;
		}

		for (std::vector<BasicBlock*>::iterator iter = bbToErase.begin();
				iter != bbToErase.end(); iter++) {
			(*iter)->eraseFromParent();
		}
		ABCD += bbToMerge.size();
		if (MY_DEBUG) {
			if (ifAllMerged) {
				errs()
						<< "Successfully remove all redundant array bound checks!"
						<< "\n\r";
			} else
				errs() << "Fail to remove all redundant array bound checks!"
						<< "\n\r";
		}
		return ifChanged;
	}
};
char MyFunctionPass::ID = 0;
RegisterPass<MyFunctionPass> X("abcd",
		"Redundant Array Bounds Checking Elimination");
}

